#!/bin/bash
set -e  # 出错立即退出

# --------------------------
# 清理残留进程 (防止占用 GPU / spawn 冲突)
# --------------------------
echo "清理残留进程..."
ps -ef | grep "spawn" | grep -v grep | awk '{print $2}' | xargs -r kill -9 || true

# --------------------------
# 配置参数
# --------------------------
NUM_GPUS=4               # 使用的 GPU 数量
TRAIN_SCRIPT="train.py"   # 训练脚本路径

# --------------------------
# 启动训练
# --------------------------
echo "开始训练任务..."
echo "使用 GPU 数量: $NUM_GPUS"
echo "使用脚本: $TRAIN_SCRIPT"

if [ "$NUM_GPUS" -eq 1 ]; then
    # ⚠️ 单 GPU 模式可能存在 bug
    echo "单 GPU 模式"
    accelerate launch $TRAIN_SCRIPT
else
    echo "多 GPU 模式"
    accelerate launch --multi_gpu --num_processes $NUM_GPUS $TRAIN_SCRIPT
fi

echo "✅ 训练完成"